#!/usr/bin/python

import string
import os

class single_fasta_sequence (object):
    """
    Assumes each line same length
    """
    def __init__ (self, 
                        file_name):
        self.file_name = file_name
        self.file = open(file_name)

        self.header_len = len(self.file.readline())
        self.line_len = len(self.file.readline()) - 1
        self.file.seek(0, os.SEEK_END)
        self.file_len = self.file.tell()
        self.file.seek(0, os.SEEK_SET)
        self.transtable = string.maketrans("","")


    def to_file_coordinates (self, pos):
        num_lines = pos / self.line_len
        return self.header_len + pos + int(num_lines)

    def get_seq(self, seq_beg, len):
        beg = self.to_file_coordinates(seq_beg)
        end = self.to_file_coordinates(seq_beg + len)
        #print beg, end, seq_beg, len
        if beg > self.file_len:
            raise "start too big: desired position %d, file_len %d" %(end, self.file_len)
        if end > self.file_len:
            raise "length too long: desired position %d, file_len %d" %(end, self.file_len)
        len = end - beg
        self.file.seek(beg, os.SEEK_SET)
        seq = self.file.read(int(len))
        return seq.translate(self.transtable, '\n')

        

        
def unit_test():
    chrm = single_fasta_sequence("test_single_fasta_sequence/mouse_chr1.fa")
    print chrm.get_seq(3050000, 50)
    chrm = single_fasta_sequence("test_single_fasta_sequence/test.fa")
    print chrm.get_seq(0, 60)
    print chrm.get_seq(30, 60)
    print chrm.get_seq(130, 60)
    print chrm.get_seq(260, 60)
if __name__ == '__main__':
    unit_test()    


